import logging
import numpy as np

from util.util_log import setup_logging
from FitnessMultiBeamTrace import FitnessMultiBeamTrace


# 配置日志，默认打印到控制台，也可以设置打印到文件
setup_logging()
# setup_logging(log_file="../files/multi-beam-trace/beam2/QE/L=L1+L4/trace_2025-03-19-1644-16beam-qe.txt")
# 获取日志记录器并记录日志
logger = logging.getLogger("[BGWO]")


class BinaryGreyWolfOptimizer:
    def __init__(self, objective_function, dimensions, population_size=30, iterations=100):
        """
        初始化二进制灰狼优化算法。

        :param objective_function: 适应度函数
        :param dimensions: 搜索空间维度 (tuple of two integers for 2D arrays)
        :param population_size: 种群大小
        :param iterations: 迭代次数
        """
        self.objective_function = objective_function
        self.dimensions = dimensions
        self.population_size = population_size
        self.iterations = iterations
        self.position = np.zeros((population_size, *dimensions))
        self.alpha_position = np.zeros(dimensions)
        self.beta_position = np.zeros(dimensions)
        self.delta_position = np.zeros(dimensions)
        self.alpha_score = float('inf')
        self.beta_score = float('inf')
        self.delta_score = float('inf')
        self.convergence_curve = []

    def __generate_population_part(self, before_array, current_array, max_num):
        part = []
        # 添加before和current作为前两个元素
        part.append(before_array.copy().astype(int))
        part.append(current_array.copy().astype(int))
        # 计算 diff 数组
        diff = before_array != current_array
        # 获取所有 False 的坐标
        false_coords = np.argwhere(diff)
        # 生成剩余的个体
        for _ in range(max_num - 2):
            # 复制 current_array 并强制转换为 int 类型
            modified_array = np.copy(current_array).astype(int)
            # 随机选择一些差异行
            if np.sum(diff == True) > 0:
                # 随机选择一部分 False 坐标
                num_to_select = np.random.randint(0, len(false_coords) + 1)
                selected_false_coords = false_coords[np.random.choice(len(false_coords), num_to_select, replace=False)]
                # 将 modified_array 中选取坐标替换为 before_array
                for coord in selected_false_coords:
                    row, col = coord
                    modified_array[row, col] = before_array[row, col]
                # print("Modified array:", modified_array.tolist())
            part.append(modified_array)
        return part

    def initialize_population(self):
        """初始化种群"""
        for i in range(self.population_size):
            self.position[i] = np.random.randint(2, size=self.dimensions)

    def initialize_population_group(self, phaseBit_mix_ps_before, phaseBit_mix_ps, phaseBit_mix_qe_before,
                                    phaseBit_mix_qe):
        parts_ps = self.__generate_population_part(phaseBit_mix_ps_before, phaseBit_mix_ps, self.population_size // 2)
        parts_qe = self.__generate_population_part(phaseBit_mix_qe_before, phaseBit_mix_qe,
                                                 self.population_size - self.population_size // 2)
        self.population = parts_ps + parts_qe

    def initialize_alpha_beta_delta(self):
        scores = [self.objective_function(pos) for pos in self.position]
        # 将位置和得分组合在一起以便排序
        positions_scores = list(zip(scores, self.position))
        # 根据得分从小到大排序
        positions_scores.sort(key=lambda x: x[0])
        # 分配 alpha, beta, delta
        self.alpha_score, self.alpha_position = positions_scores[0]
        self.beta_score, self.beta_position = positions_scores[1]
        self.delta_score, self.delta_position = positions_scores[2]

    def update_alpha_beta_delta(self, position, score, index):
        """更新 alpha, beta, delta 的位置和分数"""
        if score < self.alpha_score:
            self.delta_score = self.beta_score
            self.delta_position = self.beta_position.copy()
            self.beta_score = self.alpha_score
            self.beta_position = self.alpha_position.copy()
            self.alpha_score = score
            self.alpha_position = position.copy()
        elif score < self.beta_score:
            self.delta_score = self.beta_score
            self.delta_position = self.beta_position.copy()
            self.beta_score = score
            self.beta_position = position.copy()
        elif score < self.delta_score:
            self.delta_score = score
            self.delta_position = position.copy()

    def sigmoid(self, t):
        """Sigmoid 函数用于二进制转换"""
        return 1 / (1 + np.exp(-t))

    def optimize(self, phaseBit_mix_ps_before, phaseBit_mix_ps, phaseBit_mix_qe_before, phaseBit_mix_qe):
        """运行二进制灰狼优化算法"""

        # """初始化返回值"""
        self.best_individual = phaseBit_mix_ps
        self.best_fitness = self.objective_function(self.best_individual)
        self.best_individual_history = []
        self.best_fitness_history = []

        # """普通初始化"""
        # self.initialize_population()
        # """初始化"""
        self.initialize_population_group(phaseBit_mix_ps_before, phaseBit_mix_ps,
                                         phaseBit_mix_qe_before, phaseBit_mix_qe)

        # 初始化alpha, beta, delta
        self.initialize_alpha_beta_delta()

        for iteration in range(self.iterations):
            a = 2 - iteration * ((2) / self.iterations)  # 线性递减从2到0

            # 记录本次最佳个体及其适应度
            this_best_individual = self.alpha_position
            this_best_fitness = self.objective_function(this_best_individual)

            for i in range(self.population_size):
                A1 = a * (2 * np.random.rand() - 1)
                A2 = a * (2 * np.random.rand() - 1)
                A3 = a * (2 * np.random.rand() - 1)

                C1 = 2 * np.random.rand(*self.dimensions)
                C2 = 2 * np.random.rand(*self.dimensions)
                C3 = 2 * np.random.rand(*self.dimensions)

                D_alpha = abs(C1 * self.alpha_position - self.position[i])
                X1 = self.alpha_position - A1 * D_alpha

                D_beta = abs(C2 * self.beta_position - self.position[i])
                X2 = self.beta_position - A2 * D_beta

                D_delta = abs(C3 * self.delta_position - self.position[i])
                X3 = self.delta_position - A3 * D_delta

                # 计算新的位置
                new_position = (X1 + X2 + X3) / 3

                # 使用 Sigmoid 函数进行二进制转换
                probability = self.sigmoid(new_position)

                # 更新位置
                self.position[i] = (probability >= 0.5).astype(int)

                # 计算适应度值
                fitness = self.objective_function(self.position[i])

                # 更新 alpha, beta, delta
                self.update_alpha_beta_delta(self.position[i], fitness, i)

                # 更新本次最佳个体及其适应度
                if fitness < this_best_fitness:
                    this_best_individual = self.position[i].copy()
                    this_best_fitness = fitness

            # 记录收敛曲线
            self.convergence_curve.append(self.alpha_score)

            # 记录最佳适应度曲线
            self.best_individual = this_best_individual.copy()
            self.best_fitness = this_best_fitness
            self.best_individual_history.append(this_best_individual.copy())
            self.best_fitness_history.append(this_best_fitness)

            # 打印当前最优解
            # print(f"Iteration {iteration + 1}: Best Fitness = {self.alpha_score}")
            logger.info("iteration=%d: self.best_fitness=%f" % (iteration, self.best_fitness))
            logger.info("---------------------------------------------------------------------------")

        # return self.alpha_position, self.alpha_score, self.convergence_curve
        return self.best_individual, self.best_fitness, self.best_fitness_history, self.best_individual_history


# 示例目标函数：最大化一的数量
def maximize_ones(binary_matrix):
    return np.sum(binary_matrix)


def test():
    # 设置参数
    dimensions = (4, 4)  # Example 3x3 matrix
    population_size = 30
    iterations = 100
    # 创建一个示例 FitnessMultiBeamTrace 实例
    bit_num = 8  # 根据实际情况设置比特数
    beam_num = 2  # 根据实际情况设置波束数
    fitness_function = FitnessMultiBeamTrace(bit_num, beam_num)

    # 初始化主瓣位置
    phaseBit_target = np.random.randint(2, size=(4, 4))
    fitness_function.init_main_lope_pos(phaseBit_target)

    # 定义初始相位比特矩阵
    phaseBit_ps_before = np.array([[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]])
    phaseBit_ps = np.array([[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [1, 1, 1, 1]])
    phaseBit_qe_before = np.array([[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]])
    phaseBit_qe = np.array([[1, 1, 1, 1], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]])

    # 创建并运行二进制灰狼优化算法实例
    bgwo = BinaryGreyWolfOptimizer(objective_function=fitness_function.fitness,
                                   dimensions=dimensions,
                                   population_size=population_size,
                                   iterations=iterations)

    best_individual, best_fitness, best_fitness_history, best_individual_history \
        = bgwo.optimize(phaseBit_ps_before, phaseBit_ps, phaseBit_qe_before, phaseBit_qe)

    print(f"\nbest_individual:\n{best_individual}")
    print(f"best_fitness: {best_fitness}")


if __name__=="__main__":
    test()
